import os
from typing import Callable

import pytest

from storage.blobstore.crypto import (
    encrypt_blob_aes128_ctr, decrypt_blob_aes128_ctr,
    encrypt_blob_aes128_cbc, decrypt_blob_aes128_cbc,
    encrypt_blob_aes256_ctr, decrypt_blob_aes256_ctr,
    encrypt_blob_aes256_cbc, decrypt_blob_aes256_cbc,
)


@pytest.mark.parametrize('encrypt,decrypt,secret,nonce_iv', [
    pytest.param(encrypt_blob_aes128_ctr, decrypt_blob_aes128_ctr, os.urandom(16), os.urandom(16), id='AES128-CTR'),
    pytest.param(encrypt_blob_aes256_ctr, decrypt_blob_aes256_ctr, os.urandom(32), os.urandom(16), id='AES256-CTR'),
    pytest.param(encrypt_blob_aes128_cbc, decrypt_blob_aes128_cbc, os.urandom(16), os.urandom(16), id='AES128-CBC'),
    pytest.param(encrypt_blob_aes256_cbc, decrypt_blob_aes256_cbc, os.urandom(32), os.urandom(16), id='AES256-CBC'),
])
def test_smoking_encrypt_decrypt(encrypt: Callable, decrypt: Callable, secret: bytes, nonce_iv: bytes,
                                 plaintext_data: bytes):
    ciphertext = encrypt(secret, nonce_iv, plaintext_data)
    plaintext = decrypt(secret, nonce_iv, ciphertext)
    assert plaintext == plaintext_data


@pytest.mark.benchmark
@pytest.mark.parametrize('encrypt,secret,nonce_iv', [
    pytest.param(encrypt_blob_aes128_ctr, os.urandom(16), os.urandom(16), id='AES128-CTR'),
    pytest.param(encrypt_blob_aes256_ctr, os.urandom(32), os.urandom(16), id='AES256-CTR'),
    pytest.param(encrypt_blob_aes128_cbc, os.urandom(16), os.urandom(16), id='AES128-CBC'),
    pytest.param(encrypt_blob_aes256_cbc, os.urandom(32), os.urandom(16), id='AES256-CBC'),
])
def test_benchmark_encrypt(benchmark, encrypt: Callable, secret: bytes, nonce_iv: bytes, plaintext_data: bytes):
    benchmark(encrypt, secret, nonce_iv, plaintext_data)
